import os, sys
import numpy as np
import cv2
import shutil
import networkx as nx
cur_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(cur_dir, '../kensaku'))
from face_search import get_config
from face_search.models import get_model


def get_max_component(features, sim_thresh):
    g = nx.Graph()
    nnodes = len(features)
    for node_index in range(nnodes):
        g.add_node(node_index)
    for i in range(nnodes):
        for j in range(i, nnodes):
            feat1 = features[i]
            feat2 = features[j]
            sim = np.dot(feat1, feat2)
            if sim > sim_thresh:
                g.add_edge(i, j)
    
    # get connected components
    components = list(nx.connected_components(g))
    components_size = np.array([len(c) for c in components])
    max_index = np.argmax(components_size)

    return components[max_index]


def test_person(model, image_folder, face_folder, crop_face=False, sim_thresh=0.6):
    if not os.path.exists(face_folder):
        os.makedirs(face_folder)
    image_paths = []
    image_features = []
    for image_name in os.listdir(image_folder):
        image_path = os.path.join(image_folder, image_name)
        if not os.path.isfile(image_path): continue
        image_paths.append(image_path)
        print('processing {}'.format(image_name))
        features = model.run(image_path)
        if len(features) != 1:
            continue
        image_features.append(features[0])
    
    # use 
    indices = get_max_component(image_features, sim_thresh)
    print(indices)
    if len(indices) < 3:
        # not good face folder
        return
    
    for index in indices:
        image_path = image_paths[index]
        save_path = os.path.join(face_folder, '{:03d}.jpg'.format(index))
        if crop_face:
            crops = model.get_face_crops(image_path)
            face_crop = crops[0]
            print(face_crop.shape)
            cv2.imwrite(save_path, face_crop)
        else:
            shutil.move(image_path, save_path)


if __name__ == '__main__':
    cfg = get_config('./configs/face_search_v0.2.0.yaml')
    model = get_model(cfg.MODEL)
    image_folder = 'F:/Data/Graphis/Graphis Photo DVD09/First Gravure/2014-12-05 China Matsuoka'
    face_folder = 'F:/Projects/PornFace/data/jav_faces/China_Matsuoka'
    test_person(model, image_folder, face_folder)
